Avoid reshape allocation in extract_jacobian! for Matrix results by ChrisRackauckas-Claude · Pull Request #797 · JuliaDiff/ForwardDiff.jl

ChrisRackauckas-Claude · 2026-03-17T13:55:13Z

Summary

extract_jacobian! calls reshape(result, length(ydual), n) which allocates a 48-byte ReshapedArray wrapper on every call. Under --check-bounds=yes (which Pkg.test() always uses), this allocation cannot be elided by the compiler.

For implicit ODE/SDE solvers that call jacobian! multiple times per step via the NL solver, this adds up. For example, SKenCarp in StochasticDiffEq.jl calls nlsolve! 3 times per step, each triggering a Jacobian computation, resulting in 3 × 48 = 144 bytes/step of unnecessary allocations.

Fix: Add a specialized extract_jacobian! method for Matrix result with AbstractVector ydual that uses direct loop indexing instead of reshape+broadcast. This is zero-alloc under --check-bounds=yes and produces identical results.

MWE

julia> using ForwardDiff

julia> T = ForwardDiff.Tag{Nothing, Float64};

julia> ydual = ForwardDiff.Dual{T}.([1.0, 2.0, 3.0],
           [ForwardDiff.Partials((1.0, 2.0, 3.0)),
            ForwardDiff.Partials((4.0, 5.0, 6.0)),
            ForwardDiff.Partials((7.0, 8.0, 9.0))]);

julia> result = zeros(3, 3);

# Under --check-bounds=yes:
# Before: 48 bytes per call
# After:  0 bytes per call
julia> @allocated ForwardDiff.extract_jacobian!(T, result, ydual, 3)
0

Test plan

New allocation test in test/AllocationsTest.jl (passes under --check-bounds=yes)
New correctness test verifying extracted values match expected Jacobian
Full test suite passes (9036/9036 tests, including under --check-bounds=yes)

🤖 Generated with Claude Code

devmotion · 2026-03-17T13:58:41Z

src/jacobian.jl

+function extract_jacobian!(::Type{T}, result::Matrix, ydual::AbstractVector, n) where {T}
+    for j in 1:n
+        for i in eachindex(ydual)
+            result[i, j] = partials(T, ydual[i], j)


There's no guarantee that result has the correct dimensions? Additionally, i might not be a valid index for result.

codecov · 2026-03-17T14:03:07Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.75%. Comparing base (ff0d903) to head (ac91112).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #797   +/-   ##
=======================================
  Coverage   90.75%   90.75%           
=======================================
  Files          11       11           
  Lines        1071     1071           
=======================================
  Hits          972      972           
  Misses         99       99

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

devmotion · 2026-03-17T14:06:53Z

test/AllocationsTest.jl

+    T = ForwardDiff.Tag{Nothing, Float64}
+    N = 3
+    ydual = ForwardDiff.Dual{T}.(
+        [1.0, 2.0, 3.0],
+        [ForwardDiff.Partials((1.0, 2.0, 3.0)),
+         ForwardDiff.Partials((4.0, 5.0, 6.0)),
+         ForwardDiff.Partials((7.0, 8.0, 9.0))]
+    )
+    result = zeros(3, 3)
+
+    allocs_extract!() = @allocated ForwardDiff.extract_jacobian!(T, result, ydual, N)


Move these definitions into the function body to avoid closing over these variables?

`extract_jacobian!` called `reshape(result, length(ydual), n)` which allocates a 48-byte ReshapedArray wrapper. Under `--check-bounds=yes` (used by Pkg.test), this allocation cannot be elided by the compiler, causing 48 bytes per jacobian! call. For implicit ODE/SDE solvers that call jacobian! multiple times per step, this adds up (e.g. 144 bytes/step for SKenCarp with 3 NL solver iterations). Add `_maybe_reshape` that returns the array as-is when it already has the target shape, avoiding the wrapper allocation. Falls back to `reshape` when dimensions don't match. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas · 2026-03-17T14:30:01Z

@devmotion this should be simpler?

devmotion · 2026-03-17T15:07:02Z

src/jacobian.jl

+    if size(result) == (m, n)
+        return result
+    else
+        return reshape(result, m, n)
+    end


This is not type-stable in general?

From searching through the codebase, my impression is that in quite a few cases we call extract_jacobian! internally with a matrix of the correct dimensions constructed in the previous line. In these cases, we don't even need this check and can always operate with result.

AFAICT the only case where we might have to reshape is when users provide an output array.

Maybe we could reshape user-provided arrays to matrices higher up in the call stack, maybe even in the user-facing function directly, and then never reshape in this internal function and only accept matrices.

The reshaping of the user-provided arrays could eg be done unconditionally for non-Matrix input (and other types for which we don't know whether the function in this draft would be type-stable) and only for Matrix (and other known to be type-stable types) we would use the conditional reshaping.

Address review feedback: replace the type-unstable _maybe_reshape helper with two dispatch-based extract_jacobian! methods: - AbstractMatrix: skip reshape entirely (zero-alloc, hot path for DiffEq) - AbstractArray: reshape unconditionally (type-stable for non-matrix inputs) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas-Claude · 2026-03-24T10:44:40Z

@devmotion Good point about the type instability of _maybe_reshape. Replaced it with dispatch-based methods in 0964e46:

extract_jacobian!(::Type{T}, result::AbstractMatrix, ...) — skips reshape entirely (zero-alloc, covers the hot path where callers construct a matrix internally or users pass a Matrix)
extract_jacobian!(::Type{T}, result::AbstractArray, ...) — reshapes unconditionally (type-stable for non-matrix user-provided arrays)

Both paths are now type-stable. Full test suite passes (9035/9035) and the allocation test confirms 0 bytes for the Matrix path under --check-bounds=yes.

devmotion · 2026-03-24T10:55:02Z

src/jacobian.jl


+# Specialized method for AbstractMatrix: no reshape needed, avoids ReshapedArray allocation
+# that cannot be elided under --check-bounds=yes.
+function extract_jacobian!(::Type{T}, result::AbstractMatrix, ydual::AbstractArray, n) where {T}


Move to the function below to minimize the diff to a one-line change?

Address review: instead of a separate dispatch method for AbstractMatrix, use a conditional in the existing method to skip reshape when result is already a matrix. This minimizes the diff while preserving the allocation fix under --check-bounds=yes. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

devmotion

If you bump the version number, we can tag a release when the PR is merged.

ChrisRackauckas · 2026-03-24T12:23:36Z

done

devmotion reviewed Mar 17, 2026

View reviewed changes

ChrisRackauckas-Claude force-pushed the fix-extract-jacobian-reshape-alloc branch from cb850f0 to 526aecd Compare March 17, 2026 14:15

ChrisRackauckas-Claude force-pushed the fix-extract-jacobian-reshape-alloc branch from 526aecd to 84b9290 Compare March 17, 2026 14:29

ChrisRackauckas requested a review from devmotion March 17, 2026 14:29

ChrisRackauckas-Claude mentioned this pull request Mar 17, 2026

Tighten SKenCarp allocation test to == 0 after ForwardDiff.jl#797 SciML/StochasticDiffEq.jl#699

Closed

devmotion reviewed Mar 17, 2026

View reviewed changes

devmotion reviewed Mar 24, 2026

View reviewed changes

Update Project.toml

ac91112

devmotion approved these changes Mar 24, 2026

View reviewed changes

devmotion merged commit 1295777 into JuliaDiff:master Mar 24, 2026
37 of 52 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Avoid reshape allocation in extract_jacobian! for Matrix results#797

Avoid reshape allocation in extract_jacobian! for Matrix results#797
devmotion merged 4 commits intoJuliaDiff:masterfrom
ChrisRackauckas-Claude:fix-extract-jacobian-reshape-alloc

ChrisRackauckas-Claude commented Mar 17, 2026

Uh oh!

devmotion Mar 17, 2026

Uh oh!

codecov bot commented Mar 17, 2026 •

edited

Loading

Uh oh!

devmotion Mar 17, 2026

Uh oh!

ChrisRackauckas commented Mar 17, 2026

Uh oh!

devmotion Mar 17, 2026

Uh oh!

ChrisRackauckas-Claude commented Mar 24, 2026

Uh oh!

devmotion Mar 24, 2026

Uh oh!

devmotion left a comment

Uh oh!

ChrisRackauckas commented Mar 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

ChrisRackauckas-Claude commented Mar 17, 2026

Summary

MWE

Test plan

Uh oh!

devmotion Mar 17, 2026

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

devmotion Mar 17, 2026

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas commented Mar 17, 2026

Uh oh!

devmotion Mar 17, 2026

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas-Claude commented Mar 24, 2026

Uh oh!

devmotion Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

devmotion left a comment

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas commented Mar 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov bot commented Mar 17, 2026 •

edited

Loading